import sys
import os
import git # 没有引入执行脚本居然没报错
import importlib
import requests
import openpyxl
import re
from pathlib import Path
from deepseekclient.deepseekapi_client import DeepSeekManager

project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.append(os.path.abspath(project_root + '/client/'))
sys.path.append(os.path.abspath(project_root + '/middle/'))
uosvulapi_client = importlib.import_module('uosvulapi_client')

sys.path.append(os.path.abspath(project_root + '/deepseekclient/'))
deepseekapi_client = importlib.import_module('uosvulapi_client')


from client import config
import logging
from middle.log import setup_logging
logger = logging.getLogger(__name__)

def get_third_line(text):
    """
    爬虫使用，获取第三行
    """
    # 按行分割字符串
    lines = text.splitlines()
    # 检查是否有第三行
    if len(lines) >= 3:
        return lines[2]  # 索引从 0 开始，第三行的索引是 2
    else:
        return "没有第三行"


def get_cve_title(cve_id):
    """
    爬虫使用，根据 cveid 获取到 cve 的标题内容
    """
    url = f"https://cve.circl.lu/api/cve/{cve_id}"
    response = requests.get(url)
    if response.ok:
        cve_data = response.json()
        try:
            url = cve_data["containers"]["cna"]["references"][0]["url"]
            resolved = cve_data["containers"]["cna"]["descriptions"][0]["value"]
            title = get_third_line(resolved)
            return title
        except (KeyError, IndexError):
            pass
    return None

def get_nvd_description(cve_id):
    """
    获取 nvd 的描述内容
    """
    url = f"https://cve.circl.lu/api/cve/{cve_id}"
    response = requests.get(url)
    if response.ok:
        cve_data = response.json()
        try:
            url = cve_data["containers"]["cna"]["references"][0]["url"]
            resolved = cve_data["containers"]["cna"]["descriptions"][0]["value"]
            return resolved
        except (KeyError, IndexError):
            pass
    return None

def search_git_log(subject, branch="main", repo_path="path/to/your/repo"):
    """
    搜索 Git 仓库的 commit title（第一行）是否包含某个 subject（如 CVE ID），
    如果匹配则返回 True，否则返回 False。
    """
    logger.debug(f"subject: {subject} branch: {branch} repo: {repo_path}")
    try:
        repo = git.Repo(repo_path)
        
        for commit in repo.iter_commits(branch):
            
            # 跳过合并提交（有多个父提交）
            if len(commit.parents) > 1:
                continue
                
            commit_title = commit.message.split('\n')[0]
            
            if subject in commit_title:
                logger.info(f"匹配到 commit: {commit.hexsha} - {commit_title}")
                #return True, {
                #    'sha': commit.hexsha,
                #    'title': commit_title,
                #    'author': str(commit.author),
                #    'date': commit.authored_datetime.isoformat()
                #}
                return True
        
        logger.info("未找到匹配的非合并提交")
        return False
    except Exception as e:
        logger.error(f"发生错误: {e}")
        return False

def get_patch_content(subject, branch="main", repo_path="path/to/your/repo"):
    """
    搜索 Git 仓库的 commit title（第一行）是否包含某个 subject（如 CVE ID），
    如果匹配则返回 True，否则返回 False。
    """
    logger.debug(f"subject: {subject} branch: {branch} repo: {repo_path}")
    try:
        repo = git.Repo(repo_path)
        
        for commit in repo.iter_commits(branch):
            
            # 跳过合并提交（有多个父提交）
            if len(commit.parents) > 1:
                continue
                
            commit_title = commit.message.split('\n')[0]
            
            if subject in commit_title:
                patch_content = repo.git.diff(f'{commit.parents[0].hexsha}..{commit.hexsha}')
                return patch_content
        
        logger.info("未找到匹配的非合并提交")
        return False
    except Exception as e:
        logger.error(f"发生错误: {e}")
        return False


def get_sorted_all_data(production):
    """
    获取到排好序的所有对应产线的数据
    """
    items = []
    if production == config.Production.V25:
        items = uosvulapi_client.get_all_v25_Data()
    elif production == config.Production.V20_510E_1070:
        items = uosvulapi_client.get_all_v20_510e_1070_Data()
    elif production == config.Production.V20_510E_1060:
        items = uosvulapi_client.get_all_v20_510e_1060_Data()
    elif production == config.Production.V20_510E_1050:
        items = uosvulapi_client.get_all_v20_510e_1050_Data()
    elif production == config.Production.V20_510A_1070:
        items = uosvulapi_client.get_all_v20_510a_1070_Data()
    elif production == config.Production.V20_510A_1060:
        items = uosvulapi_client.get_all_v20_510a_1060_Data()
    elif production == config.Production.V20_510A_1050:
        items = uosvulapi_client.get_all_v20_510a_1050_Data()
    elif production == config.Production.V20_419A_1070:
        items = uosvulapi_client.get_all_v20_419A_1070_Data()
    elif production == config.Production.V20_419A_1060:
        items = uosvulapi_client.get_all_v20_419A_1060_Data()
    elif production == config.Production.V20_419A_1050:
        items = uosvulapi_client.get_all_v20_419A_1050_Data()
    elif production == config.Production.V20_419E_1070:
        items = uosvulapi_client.get_all_v20_419E_1070_Data()
    elif production == config.Production.V20_419E_1060:
        items = uosvulapi_client.get_all_v20_419E_1060_Data()
    elif production == config.Production.V20_419E_1050:
        items = uosvulapi_client.get_all_v20_419E_1050_Data()

    sorted_items = sorted(items, key=lambda x: x['cveid'])
    return sorted_items

# 为了减少本地查找数据的速度
old_cve_id = "" # 此文件的全局变量，判断是否是一组的
old_cve_update = False # 之前的同一个 cve 是不是需要更新
def update_cve_items(items, production):
    """
    更新所有的 cve 状态（核心函数）
    :param production: 这个标识的一个产线： config 当中使用的枚举值。
    多个架构的同一个名字的 cve 有不同的 cveid ，每个 cveid 都需要更新。设置一
    个 cve_group_name 是为了在同源异构的情况下，不必将多个 cve 在 git 仓库查询
    多次。判断是否需要更新。
    需要更新的数据每 500 条一组更新数据库
    """
    logger.info("进入： update_cve_items ， cared items 长度为 {}".format(len(items)))
    code_path = ""
    code_branch = ""
    fixed_version = ""
    post_data = []
    num = 0
    if production == config.Production.V25:
        code_path = config.CODE_660_PATH
        code_branch = config.CODE_660_2500_BRANCH
        fixed_version = config.FIX_660_2500_VERSION
    elif production == config.Production.V20_510E_1070:
        code_path = config.CODE_510_PATH
        code_branch = config.CODE_510_1070_BRANCH
        fixed_version = config.FIX_510E_1070_VERSION
    elif production == config.Production.V20_510E_1060:
        code_path = config.CODE_510_PATH
        code_branch = config.CODE_510_1060_BRANCH
        fixed_version = config.FIX_510E_1060_VERSION
    elif production == config.Production.V20_510E_1050:
        code_path = config.CODE_510_PATH
        code_branch = config.CODE_510_1050_BRANCH
        fixed_version = config.FIX_510E_1050_VERSION
    elif production == config.Production.V20_510A_1070:
        code_path = config.CODE_510_PATH
        code_branch = config.CODE_510_1070_BRANCH
        fixed_version = config.FIX_510A_1070_VERSION
    elif production == config.Production.V20_510A_1060:
        code_path = config.CODE_510_PATH
        code_branch = config.CODE_510_1060_BRANCH
        fixed_version = config.FIX_510A_1060_VERSION
    elif production == config.Production.V20_510A_1050:
        code_path = config.CODE_510_PATH
        code_branch = config.CODE_510_1050_BRANCH
        fixed_version = config.FIX_510A_1050_VERSION
    elif production == config.Production.V20_419A_1070:
        code_path = config.CODE_419A_PATH
        code_branch = config.CODE_419A_1070_BRANCH
        fixed_version = config.FIX_419A_1070_VERSION
    elif production == config.Production.V20_419A_1060:
        code_path = config.CODE_419A_PATH
        code_branch = config.CODE_419A_1060_BRANCH
        fixed_version = config.FIX_419A_1060_VERSION
    elif production == config.Production.V20_419A_1050:
        code_path = config.CODE_419A_PATH
        code_branch = config.CODE_419A_1050_BRANCH
        fixed_version = config.FIX_419A_1050_VERSION
    elif production == config.Production.V20_419E_1070:
        code_path = config.CODE_419E_PATH
        code_branch = config.CODE_419E_1070_BRANCH
        fixed_version = config.FIX_419E_1070_VERSION
    elif production == config.Production.V20_419E_1060:
        code_path = config.CODE_419E_PATH
        code_branch = config.CODE_419E_1060_BRANCH
        fixed_version = config.FIX_419E_1060_VERSION
    elif production == config.Production.V20_419E_1050:
        code_path = config.CODE_419E_PATH
        code_branch = config.CODE_419E_1050_BRANCH
        fixed_version = config.FIX_419E_1050_VERSION
    logger.info("code_path: {} code_branch: {} fixed_version: {}".format(code_path, code_branch
                                                                , fixed_version))
    # 发送更新的 http 请求
    for vul in items:
        global old_cve_id
        global old_cve_update
        new_cve_update = False
        new_cve_id = vul.get("cveid")
        if new_cve_id == old_cve_id:
            new_cve_update = old_cve_update
        else:
            # 通过 cveid 来获取 subject
            subject = get_cve_title(new_cve_id)
            new_cve_update = search_git_log(subject, code_branch, code_path)
            old_cve_update = new_cve_update
            old_cve_id = new_cve_id
        if new_cve_update:
            post_data.append({
                "id": vul.get("id"),
                "fixed_version": fixed_version,
                "status": "fixed",
                "fixer": config.FIXER,
                "local_patch": vul.get("local_patch"),
                "vul_solution": vul.get("vul_solution"),
                "vul_version":vul.get("vul_version")
            })
            # 关键打印
            logger.info("id: {} cveid: {} 需要更新".format(vul.get("id"), new_cve_id))
            num = num + 1
            if num == config.UPDATE_BATCH_SIZE: # 一批一批的更新
                logger.info("发送批量更新请求")
                ret = uosvulapi_client.update_multiple_items(post_data)
                if ret == 20000 :
                    post_data = []
                    num = 0
                else:
                    logger.info("位置一：批量更新失败")
                    os._exit(1)
        else:
            logger.info("id: {} cveid: {} 未修改".format(vul.get("id"), new_cve_id))
    # 不够一批的最后更新一下
    ret = uosvulapi_client.update_multiple_items(post_data)
    if ret != 20000:
        logger.info("位置二：批量更新失败")
        os._exit(1)

def filter_all_uncared_cve_items(items, production):
    """
    过滤掉所有的 vul_level 是 "low" 和 "" 的 cveid
    还有一个额外的逻辑，V20 E 版的 sw 的 items 也不关心！
    """
    cared_items = []
    for vul in items:
        vul_level = vul.get("vul_level")
        vul_arch = vul.get("architecture")
        if vul_level == "low" or vul_level == "":
            continue
        result = (production == config.Production.V20_419E_1070 or
        	production == config.Production.V20_419E_1060 or
                production == config.Production.V20_419E_1050 or
                production == config.Production.V20_510E_1070 or
                production == config.Production.V20_510E_1060 or
                production == config.Production.V20_510E_1050) and \
        	(vul_arch == "sw_64")
        if result:
            continue
        else:
            cared_items.append(vul)
    return cared_items

def update_polluted_cve_items(items, production):
    """
    :param items: 这个包含所有的需要和不需要更新的 cve，可自行判断需要更新哪些
    """
    logger.info("items 长度： {}".format(len(items)))
    # 过滤掉状态是低危和未知的 cve 号
    cared_items = filter_all_uncared_cve_items(items, production)
    # 更新状态
    update_cve_items(cared_items, production)


def get_vul_info(items, production):
    """
    打印出来不同产线需要修改的 cve 的数量和具体的 cve 号
    :param items: 这个包含所有关心的和不关心的 cve，因为后端不改接口，没办法
    我来过滤
    """
    # 过滤掉状态是低危和未知的 cve 号
    cared_items = filter_all_uncared_cve_items(items, production)
    logger.info("===============================================================")
    logger.info(f"Production: {production}")
    logger.info("===============================================================")

    no_repeat_items = []
    old_cve_id = ""
    for vul in cared_items:
        new_cve_id = vul.get("cveid")
        if new_cve_id != old_cve_id:
            no_repeat_items.append(vul)
            old_cve_id = new_cve_id
        else:
            continue

    logger.info("cve 个数： {}".format(len(no_repeat_items)))
    #for vul in no_repeat_items:
    #    logger.info(f"{cveid}")


"""
分割线，下面是打补丁用到的一些个函数
"""
def xlsx_am_process():
    logger.info("")
    patch_path = os.path.expanduser(config.AM_TEMP_PATCH_PATH)
    if not os.path.exists(patch_path):
        os.makedirs(patch_path)  # 自动创建所有缺失的父目录
        logger.info(f"目录已创建: {config.AM_TEMP_PATCH_PATH}")
    else:
        logger.info(f"目录已存在: {config.AM_TEMP_PATCH_PATH}")
        # Expand the tilde (~) to the full home directory path
    file_path = os.path.expanduser(config.AM_XLSX_PATH)
    # Load the workbook
    wb = openpyxl.load_workbook(file_path)
    sheet = wb.active  # 获取当前活动的 sheet

    # 遍历每一行
    for index, row in enumerate(sheet.iter_rows(min_row=config.XLSX_START_NUMBER+1, max_row=sheet.max_row, min_col=1, max_col=sheet.max_column), start=config.XLSX_START_NUMBER):
        logger.info(f"index value: {index}")
        first_col_value = row[0].value  # 获取第一列的值
        logger.info(f"row[0] 的值为： {first_col_value}")
    	# 假设根据 CVE 值进行修改检查，例如检查是否包含某个特定的标识
        #if first_col_value == 'CVE-2025-21684':  # 这里可以根据需要修改条件
        cvetitle = get_cve_title(first_col_value)
        row[2].value = cvetitle
        wb.save(file_path)
    	# 查一下仓库，有没有这个cve，有的话，记在表格上，写上之前已合；没有的话，进行下面的步骤。
        #has = search_git_log(cvetitle, config.AM_CODE_BRANCH, Path(config.AM_CODE_PATH).expanduser())
        #if has:
    	#    row[1].value = "之前已合"
    	#    wb.save(file_path)
    	#    logger.info(f"{index} 之前已合")
    	#    continue

        # 根据 cvetitle 拿到 commit 号
        #logger.info(f"Path: {Path(config.UPSTREAM_CODE_PATH).expanduser()}")
        commit = get_commit(cvetitle,
                            config.UPSTREAM_CODE_BRANCH,
                            Path(config.UPSTREAM_CODE_PATH).expanduser()
                            )
        logger.info(f"commit: {commit}")
        # 可能存在一种情况，仓库分支中没有找到这个补丁，commit 就是 False，这种情况记录到表格里面去
        if commit is False:
    	    row[1].value = '无此补丁'
    	    wb.save(file_path)
    	    continue

        # 通过 commit 号生成 patch
        format_name = generate_patch_with_commitid(commit,
                                                   Path(config.UPSTREAM_CODE_PATH).expanduser(),
                                                   index)
        logger.info(f"format_name: {format_name}")
        modify_patch_subject(format_name, first_col_value)
        # 尝试打 patch ，首先把目标仓库做好清理。不能打上的要记录到表当中。
        amflag = apply_patch_with_abort(Path(config.AM_CODE_PATH).expanduser(),
                                        Path(config.AM_TEMP_PATCH_PATH).expanduser() / format_name)
        if amflag[0] is True:
            # 记录表格成功打上并删除补丁
            row[1].value = '成功打上'
            #os.remove(format_name)
        else:
            # 记录表格没有成功打上并保留补丁
            row[1].value = '未应用'
        wb.save(file_path)
        # 未打上的补丁开幕爆击，一个补丁包含 CVE 补丁中的信息，所以要如何，优化的内容是只
        # 对比 subject 的 title 而不对比整个 message

"""
对比 subject 的 title 且要是非合并提交，这个函数做了改进。
"""
def get_commit(subject, branch="main", repo_path="path/to/your/repo"):
    """
    搜索 Git 仓库中第一个非合并提交且 commit title 包含 subject 的提交
    
    Args:
        subject: 要搜索的字符串（如 CVE ID）
        branch: 要搜索的分支（默认 main）
        repo_path: 仓库路径
        max_commits: 最多检查多少个提交（None 表示不限制）
    
    Returns:
        bool: 是否找到匹配的提交
        dict/None: 如果找到返回提交信息，否则返回 None
    """
    logger.debug(f"subject: {subject} branch: {branch} repo: {repo_path}")
    try:
        repo = git.Repo(repo_path)
        
        for commit in repo.iter_commits(branch):
            
            # 跳过合并提交（有多个父提交）
            if len(commit.parents) > 1:
                continue
                
            commit_title = commit.message.split('\n')[0]
            
            if subject in commit_title:
                logger.info(f"匹配到 commit: {commit.hexsha} - {commit_title}")
                #return True, {
                #    'sha': commit.hexsha,
                #    'title': commit_title,
                #    'author': str(commit.author),
                #    'date': commit.authored_datetime.isoformat()
                #}
                return commit.hexsha
        
        logger.info("未找到匹配的非合并提交")
        return False
    except Exception as e:
        logger.error(f"发生错误: {e}")
        return False

def generate_patch_with_commitid(commit_id, repo_path="path/to/your/repo", startnumber=1):
    try:
        repo = git.Repo(repo_path)
        if not repo.bare:
            commit = repo.commit(commit_id)

            patch_num = f"{startnumber:04d}"
            subject = re.sub(r'[^\w-]', '_', commit.message.splitlines()[0].strip())
            filename = Path(config.AM_TEMP_PATCH_PATH).expanduser() / f"{patch_num}-{subject}.patch"

            # 避免覆盖已存在文件
            # if os.path.exists(filename):
            #     raise FileExistsError(f"Patch file {filename} already exists")

            with open(filename, 'w') as f:
                f.write(repo.git.format_patch('-1', '--stdout', commit.hexsha))

            return filename
        else:
            raise ValueError("Not a valid git repository")
    except Exception as e:
        print(f"Error generating patch: {str(e)}")
        return None


def apply_patch_with_abort(repo_path, patch_file):
    """
    使用 GitPython 应用 patch 并自动处理冲突

    :param repo_path: Git 仓库路径
    :param patch_file: patch 文件路径
    :return: (success: bool, message: str)
    """
    repo = git.Repo(repo_path)

    try:
        # 使用 git.am 应用 patch
        with repo.git.custom_environment(
            GIT_AUTHOR_NAME=config.AM_GIT_NAME,
            GIT_AUTHOR_EMAIL=config.AM_GIT_EMAIL
        ):
            repo.git.am("-s", patch_file)
        logger.info(f"{patch_file} 成功打上")
        return True, f"应用成功"

    except Exception as e:
        # 发生冲突时执行 abort
        logger.info(f"{patch_file} 未应用")
        try:
            repo.git.am('--abort')
            return False, f"应用失败并已回滚: {str(e)}"
        except Exception as abort_error:
            return False, f"应用失败且回滚失败: {str(e)} | 回滚错误: {str(abort_error)}"

"""
修改 patch 原来的格式，并且使用 git am -s 加上自己的签名。
"""
def modify_patch_subject(patch_file, cve_id):
    with open(patch_file, 'r+') as f:
        content = f.read()
        # 在 [PATCH] 后添加 CVE 信息
        modified = re.sub(
            r'^(Subject: \[PATCH\])',
            fr'\1 mainline: {{{cve_id}}}',
            content,
            flags=re.MULTILINE
        )
        # 写回文件
        f.seek(0)
        f.write(modified)
        f.truncate()

"""
deepseek 使用的方法
"""
def get_cve_exploit(cveid, des, title):
    """
    获取一个 cve 的攻击方法。
    @title: 可选
    @des： 可选
    """
    deepseek = DeepSeekManager()
    cve_exploit_prompt = """
你是一个资深Linux内核漏洞分析师。请严格按照以下模板分析CVE漏洞：

**CVE-ID**: [漏洞编号]  
**漏洞类型**: [如UAF/OOB/整数溢出]  
**触发流程**（分步骤描述关键函数和操作）:  
1. [触发入口]: [调用链起点，如系统调用/用户输入]  
2. [关键函数1]: [函数名及作用]  
3. [关键操作]: [如内存分配/释放/复制等]  
4. [漏洞点]: [具体错误操作，如use-after-free]  
5. [影响]: [崩溃/权限提升/信息泄露]  

**示例输出格式**（以CVE-2022-49842为例）:  
**CVE-ID**: CVE-2022-49842  
**漏洞类型**: Use-After-Free  
**触发流程**:  
1. 模块卸载: 用户执行`rmmod`命令  
2. 驱动清理: `snd_soc_exit()`调用`snd_soc_util_exit()`  
3. 重复释放: `platform_device_unregister(soc_dummy_dev)`对已释放设备二次操作  
4. 漏洞触发: 内核访问无效内存地址`ffff888008655050`  
5. 影响: 内核崩溃或潜在权限提升  

请按此格式分析以下CVE：[输入你的CVE编号或描述]
另外除了格式中要求的内容之外，不要添加任何其它无关的回答。
    """
    deepseek.add_wrapper("cve_exploit_prompt", config.DEEPSEEK_API_KEY,
                         config.DEEPSEEK_API_URL, cve_exploit_prompt, False)

    if des is None:
        des = get_nvd_description(cveid)
        # logger.info(f"des： {des}")
        # 使用上游的 master 分支
        title = get_third_line(des)
    patch_content = get_patch_content(title, "master", config.UPSTREAM_CODE_PATH)
    question = f"输出 {cveid} 的触发流程。\n\n该漏洞的NVD描述如下：\n{des}\n\n这个 CVE 补丁的内容如下：\n{patch_content}"
    logger.info(f"问题是：{question}")

    answer = deepseek.get_wrapper("cve_exploit_prompt").ask_question(question)
    logger.info(f"回答是： {answer}")
    return answer

def xlsx_get_exploits(xlsx_path):
    """
    获取表格中一系列的攻击方法。并写入表格。
    """
    #get_cve_exploit("CVE-2025-40325")
    file_path = os.path.expanduser(xlsx_path)
    
    wb = openpyxl.load_workbook(file_path)
    sheet = wb.active  # 获取当前活动的 sheet

    # 遍历每一行
    for index, row in enumerate(sheet.iter_rows(min_row=config.XLSX_START_NUMBER+1, max_row=sheet.max_row, min_col=1, max_col=sheet.max_column), start=config.XLSX_START_NUMBER):
        logger.info(f"index value: {index}")
        cveid = row[0].value  # 获取第一列的值
        logger.info(f"row[0] 的值为： {cveid}")
        des = get_nvd_description(cveid)
        cvetitle = get_third_line(des)
        row[2].value = cvetitle
        wb.save(file_path)
    	
        commit = get_commit(cvetitle,
                            "master",
                            Path(config.UPSTREAM_CODE_PATH).expanduser()
                            )
        logger.info(f"commit: {commit}")
        # 可能存在一种情况，仓库分支中没有找到这个补丁，commit 就是 False，这种情况记录到表格里面去
        if commit is False:
    	    row[1].value = '无此补丁'
    	    wb.save(file_path)
    	    continue
        else:
            row[1].value = '有此补丁'
            wb.save(file_path)

        exploit = get_cve_exploit(cveid, des, cvetitle)
        row[3].value = exploit
        wb.save(file_path)
